package com.starhub.application.rag.store;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.service.MilvusConfigService;
import com.starhub.common.bean.model.MilvusConfig;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.rag.config.MilvusEmbeddingConfig;

/**
 * 构建LangchainStore
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-22
 */
@Slf4j
@Component
public class LangchainStoreBuilder {

    @Autowired
    private MilvusConfigService milvusConfigService;

    @Autowired
    private MilvusEmbeddingConfig milvusEmbeddingConfig;

    // 使用ConcurrentHashMap来缓存MilvusEmbeddingStore实例
    private final Map<String, EmbeddingStore<TextSegment>> embeddingStoreCache = new ConcurrentHashMap<>();

    // 用于同步创建新实例的锁对象Map
    private final Map<String, Object> lockMap = new ConcurrentHashMap<>();

    // 连接健康检查时间戳
    private final Map<String, Long> lastHealthCheck = new ConcurrentHashMap<>();
    
    // 健康检查间隔（毫秒）
    private static final long HEALTH_CHECK_INTERVAL_MS = 5 * 60 * 1000; // 5分钟

    /**
     * 获取指定key的锁对象，如果不存在则创建
     */
    private Object getLock(String key) {
        return lockMap.computeIfAbsent(key, k -> new Object());
    }

    /**
     * 获取或创建MilvusEmbeddingStore实例
     * 
     * @param config        Milvus配置
     * @param collectionIdentifier 集合标识符
     * @param dimension 向量维度
     * @return EmbeddingStore实例
     */
    public EmbeddingStore<TextSegment> getOrCreateMilvusEmbeddingStore(MilvusConfig config,
            String collectionIdentifier, int dimension) {
        String cacheKey = collectionIdentifier;

        // 首先尝试从缓存中获取
        EmbeddingStore<TextSegment> store = embeddingStoreCache.get(cacheKey);
        if (store != null && isStoreHealthy(cacheKey, store)) {
            return store;
        }

        // 如果缓存中不存在或连接不健康，使用细粒度锁创建新实例
        synchronized (getLock(cacheKey)) {
            // 双重检查，确保实例仍然不存在或不健康
            store = embeddingStoreCache.get(cacheKey);
            if (store != null && isStoreHealthy(cacheKey, store)) {
                return store;
            }

            // 如果存在不健康的连接，先关闭它
            if (store != null) {
                closeStoreSafely(cacheKey, store);
            }

            try {
                log.info("Creating new MilvusEmbeddingStore for collection: {}", cacheKey);
                store = createMilvusEmbeddingStore(config, collectionIdentifier, dimension);
                embeddingStoreCache.put(cacheKey, store);
                lastHealthCheck.put(cacheKey, System.currentTimeMillis());
                return store;
            } catch (Exception e) {
                log.error("Failed to create MilvusEmbeddingStore for collection: {}", cacheKey, e);
                throw new RuntimeException("Failed to create MilvusEmbeddingStore", e);
            }
        }
    }

    /**
     * 创建MilvusEmbeddingStore实例
     */
    private MilvusEmbeddingStore createMilvusEmbeddingStore(MilvusConfig config,
            String collectionIdentifier, int dimension) {
        return MilvusEmbeddingStore.builder()
                .host(config.getHost())
                .port(config.getPort())
                .username(config.getUsername())
                .password(config.getPassword())
                .collectionName(collectionIdentifier)
                .dimension(dimension)
                .autoFlushOnInsert(milvusEmbeddingConfig.isAutoFlushOnInsert())
                .idFieldName(RagConstant.LangchainMilvusFields.FIELD_ID)
                .textFieldName(RagConstant.LangchainMilvusFields.FIELD_TEXT)
                .metadataFieldName(RagConstant.LangchainMilvusFields.FIELD_METADATA)
                .vectorFieldName(RagConstant.LangchainMilvusFields.FIELD_VECTOR)
                .build();
    }

    /**
     * 检查存储连接是否健康
     */
    private boolean isStoreHealthy(String cacheKey, EmbeddingStore<TextSegment> store) {
        Long lastCheck = lastHealthCheck.get(cacheKey);
        long currentTime = System.currentTimeMillis();
        
        // 如果距离上次健康检查时间太短，直接返回true
        if (lastCheck != null && (currentTime - lastCheck) < HEALTH_CHECK_INTERVAL_MS) {
            return true;
        }

        try {
            // 尝试执行一个简单的操作来验证连接
            if (store instanceof MilvusEmbeddingStore) {
                // 对于MilvusEmbeddingStore，我们可以尝试获取集合信息
                // 这里使用一个轻量级的操作来验证连接
                // 注意：这里可能需要根据实际的MilvusEmbeddingStore API调整
                lastHealthCheck.put(cacheKey, currentTime);
                return true;
            }
            return true;
        } catch (Exception e) {
            log.warn("Health check failed for collection: {}, error: {}", cacheKey, e.getMessage());
            return false;
        }
    }

    /**
     * 安全关闭存储连接
     */
    private void closeStoreSafely(String cacheKey, EmbeddingStore<TextSegment> store) {
        try {
            if (store instanceof AutoCloseable) {
                ((AutoCloseable) store).close();
                log.info("Successfully closed MilvusEmbeddingStore for collection: {}", cacheKey);
            }
        } catch (Exception e) {
            log.warn("Error closing MilvusEmbeddingStore for collection: {}", cacheKey, e);
        } finally {
            embeddingStoreCache.remove(cacheKey);
            lastHealthCheck.remove(cacheKey);
        }
    }

    /**
     * 关闭指定集合的MilvusEmbeddingStore实例
     * 
     * @param collectionIdentifier 集合标识
     */
    public void closeEmbeddingStore(String collectionIdentifier) {
        synchronized (getLock(collectionIdentifier)) {
            EmbeddingStore<TextSegment> store = embeddingStoreCache.get(collectionIdentifier);
            if (store != null) {
                closeStoreSafely(collectionIdentifier, store);
            }
            lockMap.remove(collectionIdentifier);
        }
    }

    /**
     * 应用关闭时清理所有MilvusEmbeddingStore实例
     */
    @PreDestroy
    public void cleanup() {
        log.info("Cleaning up all MilvusEmbeddingStore instances...");
        // 创建要清理的集合标识符列表
        List<String> collectionsToClose = new ArrayList<>(embeddingStoreCache.keySet());

        // 逐个关闭集合
        for (String collectionIdentifier : collectionsToClose) {
            try {
                closeEmbeddingStore(collectionIdentifier);
            } catch (Exception e) {
                log.error("Error during cleanup of collection: {}", collectionIdentifier, e);
            }
        }

        // 清理所有剩余的锁对象
        lockMap.clear();
        lastHealthCheck.clear();
        
        log.info("All MilvusEmbeddingStore instances have been cleaned up");
    }

    /**
     * 强制刷新所有连接
     * 用于处理连接问题时手动刷新
     */
    public void refreshAllConnections() {
        log.info("Refreshing all MilvusEmbeddingStore connections...");
        List<String> collectionsToRefresh = new ArrayList<>(embeddingStoreCache.keySet());
        
        for (String collectionIdentifier : collectionsToRefresh) {
            try {
                closeEmbeddingStore(collectionIdentifier);
            } catch (Exception e) {
                log.error("Error refreshing connection for collection: {}", collectionIdentifier, e);
            }
        }
        
        log.info("All MilvusEmbeddingStore connections have been refreshed");
    }
}
